from fastapi import APIRouter, HTTPException, File, Form, UploadFile
from pydantic import BaseModel
import os
from fastapi.responses import JSONResponse, HTMLResponse, FileResponse
import random
import string
import time
import uuid
from datetime import datetime
from typing import Dict, List, Optional
from config import redis_url, redis_host, redis_port, redis_db, arq_queue_name, upload_dir, static_files_url, app_host, app_port, redis_conn_timeout, redis_conn_retry
from arq.connections import RedisSettings, create_pool
from arq.jobs import Job
from worker import task_status, task_paused, update_total_progress

router_upload_real = APIRouter(
    tags=["ARQ异步上传"],
    responses={404: {"description": "Not found"}},
)

# 异步Redis连接池，将在路由器中使用
redis_pool = None
redis_connected = False

UPLOAD_DIR = os.path.join(os.path.dirname(__file__), upload_dir)
os.makedirs(UPLOAD_DIR, exist_ok=True)

@router_upload_real.on_event("startup")
async def startup_event():
    """初始化Redis连接池"""
    global redis_pool, redis_connected
    try:
        redis_settings = RedisSettings(
            host=redis_host,
            port=redis_port,
            database=redis_db,
            conn_timeout=redis_conn_timeout,
            conn_retries=redis_conn_retry,
            conn_retry_delay=1.0
        )
        redis_pool = await create_pool(redis_settings)
        redis_connected = True
        print(f"ARQ Redis连接池已初始化")
    except Exception as e:
        print(f"无法连接到Redis服务器: {str(e)}")
        print("应用程序将以离线模式运行，队列功能将不可用")
        redis_connected = False

@router_upload_real.on_event("shutdown")
async def shutdown_event():
    """关闭Redis连接池"""
    global redis_pool
    if redis_pool:
        try:
            # 关闭ArqRedis连接池
            redis_pool.close()
            print("ARQ Redis连接池已关闭")
        except Exception as e:
            print(f"关闭ARQ Redis连接池时出错: {str(e)}")


@router_upload_real.post("/upload/")
async def upload_files(files: List[UploadFile] = File(...)):
    """上传多个文件并添加到处理队列"""
    global redis_connected
    
    if not files:
        raise HTTPException(status_code=400, detail="没有文件上传")
    
    # 创建新的任务ID
    task_id = str(uuid.uuid4())
    
    # 初始化任务状态
    task_status[task_id] = {
        "id": task_id,
        "status": "pending",
        "progress": 0,
        "total_files": len(files),
        "created_at": datetime.now().isoformat(),
        "files": {}
    }
    
    # 初始化暂停状态
    task_paused[task_id] = False
    
    # 处理每个上传的文件
    for file in files:
        try:
            # 生成文件ID和保存路径
            file_id = str(uuid.uuid4())
            file_extension = os.path.splitext(file.filename)[1] if file.filename else ""
            save_filename = f"{file_id}{file_extension}"
            file_path = os.path.join(UPLOAD_DIR, save_filename)
            
            # 将文件保存到磁盘
            file_content = await file.read()
            with open(file_path, "wb") as f:
                f.write(file_content)
            
            # 获取文件大小
            file_size = os.path.getsize(file_path)
            
            # 生成文件URL
            file_url = f"http://{app_host}:{app_port}{static_files_url}/uploads/{save_filename}"
            
            # 记录文件信息
            task_status[task_id]["files"][file_id] = {
                "id": file_id,
                "original_filename": file.filename,
                "saved_filename": save_filename,
                "path": file_path,
                "size": file_size,
                "status": "queued" if redis_connected else "completed",  # 如果Redis未连接，直接设为完成
                "progress": 0 if redis_connected else 100,  # 如果Redis未连接，进度为100%
                "queued_at": datetime.now().isoformat(),
                "file_url": file_url
            }
            
            print(f"文件已上传: {file.filename}, ID: {file_id}, URL: {file_url}")
            
            # 如果Redis已连接，将文件处理任务添加到ARQ队列
            if redis_connected:
                try:
                    job = await redis_pool.enqueue_job(
                        "process_file",
                        file_path,
                        file_id,
                        task_id,
                        file.filename,
                        file_size,
                        _job_id=file_id
                    )
                    
                    # 记录任务ID
                    task_status[task_id]["files"][file_id]["job_id"] = job.job_id
                    
                    # 立即将状态更新为已完成，因为文件已经上传完成
                    task_status[task_id]["files"][file_id].update({
                        "status": "completed",
                        "progress": 100,
                        "completed_at": datetime.now().isoformat()
                    })
                except Exception as e:
                    print(f"将任务添加到队列时出错: {str(e)}")
                    task_status[task_id]["files"][file_id]["status"] = "error"
                    task_status[task_id]["files"][file_id]["error"] = f"队列错误: {str(e)}"
            else:
                # Redis未连接，设置状态为已完成(模拟模式)
                task_status[task_id]["files"][file_id].update({
                    "status": "completed",
                    "progress": 100,
                    "completed_at": datetime.now().isoformat(),
                    "note": "离线模式 - 队列功能不可用"
                })
            
        except Exception as e:
            # 记录错误信息
            if "file_id" in locals() and file_id in task_status[task_id]["files"]:
                task_status[task_id]["files"][file_id].update({
                    "status": "failed",
                    "error": str(e)
                })
            print(f"上传文件失败: {str(e)}")
    
    # 更新总进度
    update_total_progress(task_id)
    
    # 设置任务状态为已完成，因为文件处理实际上已经完成
    task_status[task_id]["status"] = "completed"
    task_status[task_id]["completed_at"] = datetime.now().isoformat()
    
    return {
        "task_id": task_id,
        "message": f"已成功将{len(files)}个文件添加到处理队列" + ("" if redis_connected else " (离线模式)"),
        "status_url": f"/router_upload/status/{task_id}",
        "queue_status": "online" if redis_connected else "offline"
    }


@router_upload_real.get("/status/{task_id}")
def get_task_status(task_id: str):
    """获取任务状态"""
    if task_id not in task_status:
        raise HTTPException(status_code=404, detail="任务不存在")
    
    # 添加暂停状态信息
    task_info = task_status[task_id].copy()
    task_info["paused"] = task_paused.get(task_id, False)
    
    # 添加队列连接状态
    task_info["queue_status"] = "online" if redis_connected else "offline"
    
    return task_info


@router_upload_real.post("/pause/{task_id}")
async def pause_task(task_id: str):
    """暂停任务处理"""
    global redis_connected
    
    if task_id not in task_status:
        raise HTTPException(status_code=404, detail="任务不存在")
    
    if not redis_connected:
        return {"success": False, "message": "队列服务离线，无法暂停任务"}
    
    if task_status[task_id]["status"] == "completed":
        return {"success": False, "message": "任务已完成，无法暂停"}
    
    if task_status[task_id]["status"] == "failed":
        return {"success": False, "message": "任务已失败，无法暂停"}
    
    # 设置暂停标志
    task_paused[task_id] = True
    
    # 更新任务状态
    task_status[task_id]["status"] = "paused"
    task_status[task_id]["paused_at"] = datetime.now().isoformat()
    
    # 尝试取消队列中的作业
    try:
        for file_id, file_info in task_status[task_id]["files"].items():
            if file_info["status"] == "queued":
                try:
                    job = Job(file_info["job_id"], redis_pool)
                    await job.abort()
                    file_info["status"] = "paused"
                except Exception:
                    pass
    except Exception as e:
        print(f"暂停任务时出错: {str(e)}")
    
    return {"success": True, "message": "任务已暂停"}


@router_upload_real.post("/resume/{task_id}")
async def resume_task(task_id: str):
    """恢复暂停的任务"""
    global redis_connected
    
    if task_id not in task_status:
        raise HTTPException(status_code=404, detail="任务不存在")
    
    if not redis_connected:
        return {"success": False, "message": "队列服务离线，无法恢复任务"}
    
    if not task_paused.get(task_id, False):
        return {"success": False, "message": "任务未处于暂停状态"}
    
    # 取消暂停标志
    task_paused[task_id] = False
    
    # 更新任务状态
    task_status[task_id]["status"] = "processing"
    task_status[task_id]["resumed_at"] = datetime.now().isoformat()
    
    # 重新入队暂停的文件
    try:
        for file_id, file_info in task_status[task_id]["files"].items():
            if file_info["status"] == "paused":
                # 将文件重新加入队列
                job = await redis_pool.enqueue_job(
                    "process_file",
                    file_info["path"],
                    file_id,
                    task_id,
                    file_info["original_filename"],
                    file_info["size"],
                    _job_id=file_id
                )
                
                # 更新状态
                file_info["status"] = "queued"
                file_info["job_id"] = job.job_id
    except Exception as e:
        print(f"恢复任务时出错: {str(e)}")
        return {"success": False, "message": f"恢复任务时出错: {str(e)}"}
    
    return {"success": True, "message": "任务已恢复"}


@router_upload_real.delete("/task/{task_id}")
async def delete_task(task_id: str):
    """删除任务"""
    global redis_connected
    
    if task_id not in task_status:
        raise HTTPException(status_code=404, detail="任务不存在")
    
    # 尝试取消队列中的作业
    if redis_connected:
        try:
            for file_id, file_info in task_status[task_id]["files"].items():
                try:
                    if "status" in file_info and file_info["status"] in ["queued", "processing"] and "job_id" in file_info:
                        try:
                            job = Job(file_info["job_id"], redis_pool)
                            await job.abort()
                        except Exception as e:
                            print(f"取消任务失败: {str(e)}")
                except KeyError as e:
                    print(f"删除任务时出现KeyError: {str(e)}, file_id={file_id}")
                    continue
        except Exception as e:
            print(f"取消任务时出错: {str(e)}")
    
    # 尝试删除上传的文件
    for file_id, file_info in task_status[task_id]["files"].items():
        try:
            if "path" in file_info and os.path.exists(file_info["path"]):
                os.remove(file_info["path"])
        except Exception as e:
            print(f"删除文件失败: {str(e)}")
    
    # 删除任务状态
    del task_status[task_id]
    if task_id in task_paused:
        del task_paused[task_id]
    
    return {"success": True, "message": "任务已删除"}


@router_upload_real.delete("/file/{task_id}/{file_id}")
async def delete_file(task_id: str, file_id: str):
    """删除特定文件任务"""
    global redis_connected
    
    if task_id not in task_status:
        raise HTTPException(status_code=404, detail="任务不存在")
    
    if file_id not in task_status[task_id]["files"]:
        raise HTTPException(status_code=404, detail="文件不存在")
    
    file_info = task_status[task_id]["files"][file_id]
    
    # 尝试取消队列中的作业
    if redis_connected and file_info["status"] in ["queued", "processing"]:
        try:
            job = Job(file_info["job_id"], redis_pool)
            await job.abort()
        except Exception as e:
            print(f"取消文件任务时出错: {str(e)}")
    
    # 尝试删除上传的文件
    try:
        if os.path.exists(file_info["path"]):
            os.remove(file_info["path"])
            print(f"文件已删除: {file_info['path']}")
        else:
            print(f"文件不存在: {file_info['path']}")
    except Exception as e:
        print(f"删除文件失败: {str(e)}")
    
    # 从任务中删除文件
    del task_status[task_id]["files"][file_id]
    
    # 如果任务中没有文件了，删除整个任务
    if not task_status[task_id]["files"]:
        del task_status[task_id]
        if task_id in task_paused:
            del task_paused[task_id]
        return {"success": True, "message": "文件已删除，任务已清空"}
    
    # 更新任务进度
    update_total_progress(task_id)
    
    return {"success": True, "message": "文件已删除"}


@router_upload_real.get("/download/{task_id}/{file_id}")
def get_file_download(task_id: str, file_id: str):
    """获取文件直接下载链接"""
    if task_id not in task_status:
        raise HTTPException(status_code=404, detail="任务不存在")
    
    if file_id not in task_status[task_id]["files"]:
        raise HTTPException(status_code=404, detail="文件不存在")
    
    file_info = task_status[task_id]["files"][file_id]
    
    # 检查文件是否存在
    if not os.path.exists(file_info["path"]):
        raise HTTPException(status_code=404, detail="文件不存在于服务器上")
    
    return FileResponse(
        path=file_info["path"], 
        filename=file_info["original_filename"],
        media_type="application/octet-stream"
    )


@router_upload_real.get("/upload-page")
def get_upload_page():
    """返回上传页面"""
    upload_page_path = os.path.join(os.path.dirname(__file__), "static", "upload_client.html")
    return FileResponse(upload_page_path)


@router_upload_real.get("/queue-status")
def get_queue_status():
    """获取队列状态信息"""
    return {
        "connected": redis_connected,
        "status": "online" if redis_connected else "offline",
        "queue_name": arq_queue_name,
        "tasks_count": len(task_status),
        "files_count": sum(len(task["files"]) for task in task_status.values())
    }


@router_upload_real.get("/")
async def hello():
    """测试接口"""
    return {"message": "Hello World from module"}